"""
a06_export_onnx.py - 音频分类模型导出为ONNX格式的优化脚本

功能：
1. 将训练好的PyTorch模型导出为ONNX格式
2. 支持动态批次维度
3. 添加详细的输入输出说明
4. 包含模型验证步骤
5. 优化后的日志记录

使用方法：
基本用法  python a06_export_onnx.py --model_path models/trained/model.pth --output_path models/exported/model.onnx --num_classes 2
自定义输入尺寸 python a06_export_onnx.py --model_path models/trained/model.pth --output_path models/exported/model.onnx --num_classes 2 --input_height 256 --input_width 256
"""

import torch
import torch.onnx
import argparse
import logging
from pathlib import Path
from typing import Tuple
# 导入自定义模型 (确保在PYTHONPATH中或使用正确路径)
from utils.audio_cnn import AudioCNN

# 设置日志
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def load_model(model_path: str, num_classes: int) -> torch.nn.Module:
    """
    加载训练好的PyTorch模型

    参数:
        model_path: 模型权重文件路径
        num_classes: 类别数量

    返回:
        加载好的PyTorch模型
    """
    try:
        # 初始化模型结构
        model = AudioCNN(num_classes=num_classes)

        # 加载模型权重 (支持多GPU训练保存的模型)
        state_dict = torch.load(model_path, map_location='cpu')

        # 处理多GPU训练保存的模型 (移除'module.'前缀)
        if all(k.startswith('module.') for k in state_dict.keys()):
            state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}

        model.load_state_dict(state_dict)
        model.eval()

        logger.info(f"成功加载模型: {model_path}")
        return model
    except Exception as e:
        logger.error(f"加载模型失败: {str(e)}")
        raise


def validate_model(model: torch.nn.Module, input_shape: Tuple[int, int, int, int]) -> None:
    """
    验证模型是否能正确处理输入

    参数:
        model: 要验证的模型
        input_shape: 输入张量形状 (batch, channel, height, width)
    """
    try:
        dummy_input = torch.randn(input_shape)
        output = model(dummy_input)

        logger.info(f"模型验证成功 - 输入形状: {input_shape} 输出形状: {output.shape}")
        logger.info(f"示例输出: {output[0].detach().numpy()}")
    except Exception as e:
        logger.error(f"模型验证失败: {str(e)}")
        raise


def export_to_onnx(
        model: torch.nn.Module,
        output_path: str,
        input_shape: Tuple[int, int, int, int] = (1, 1, 128, 128),
        opset_version: int = 13
) -> None:
    """
    将PyTorch模型导出为ONNX格式

    参数:
        model: 要导出的PyTorch模型
        output_path: ONNX模型保存路径
        input_shape: 输入张量形状 (batch, channel, height, width)
        opset_version: ONNX算子集版本
    """
    try:
        # 准备输出目录
        output_path = Path(output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # 创建虚拟输入
        dummy_input = torch.randn(input_shape)

        # 导出模型
        torch.onnx.export(
            model,
            dummy_input,
            str(output_path),
            export_params=True,
            opset_version=opset_version,
            do_constant_folding=True,
            input_names=["mel_spectrogram"],
            output_names=["class_probabilities"],
            dynamic_axes={
                "mel_spectrogram": {0: "batch_size"},
                "class_probabilities": {0: "batch_size"}
            },
            verbose=False
        )

        logger.info(f"模型成功导出到: {output_path}")

        # 验证导出的ONNX模型
        import onnx
        onnx_model = onnx.load(str(output_path))
        onnx.checker.check_model(onnx_model)
        logger.info("ONNX模型验证通过")

        # 打印模型信息
        logger.info(f"输入节点: {onnx_model.graph.input[0]}")
        logger.info(f"输出节点: {onnx_model.graph.output[0]}")

    except Exception as e:
        logger.error(f"导出ONNX模型失败: {str(e)}")
        raise


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='导出PyTorch模型到ONNX格式')
    parser.add_argument('--model_path', type=str, required=True,
                        help='PyTorch模型权重路径')
    parser.add_argument('--output_path', type=str, required=True,
                        help='ONNX模型输出路径')
    parser.add_argument('--num_classes', type=int, required=True,
                        help='模型类别数量')
    parser.add_argument('--input_height', type=int, default=128,
                        help='输入特征高度（默认为128）')
    parser.add_argument('--input_width', type=int, default=128,
                        help='输入特征宽度（默认为128）')
    return parser.parse_args()


def main():
    # 解析参数
    args = parse_args()

    # 加载模型
    model = load_model(args.model_path, args.num_classes)

    # 定义输入形状 (batch, channel, height, width)
    input_shape = (1, 1, args.input_height, args.input_width)

    # 验证模型
    validate_model(model, input_shape)

    # 导出ONNX模型
    export_to_onnx(
        model=model,
        output_path=args.output_path,
        input_shape=input_shape
    )


if __name__ == "__main__":


    main()